{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a Semantic Segmentation Model with TIMM Encoders\n",
    "\n",
    "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/geoai/blob/main/docs/examples/train_timm_segmentation.ipynb)\n",
    "\n",
    "This notebook demonstrates how to train semantic segmentation models using [PyTorch Image Models (timm)](https://github.com/huggingface/pytorch-image-models) encoders. This approach combines:\n",
    "\n",
    "- **1000+ TIMM Encoders**: State-of-the-art backbones (ResNet, EfficientNet, ViT, ConvNeXt, etc.)\n",
    "- **9 Architectures**: U-Net, U-Net++, DeepLabV3+, FPN, PSPNet, LinkNet, MANet, PAN\n",
    "- **Multi-channel Support**: RGB, RGBN, or any number of input channels\n",
    "- **Simplified API**: Similar to `train_segmentation_model` for ease of use\n",
    "\n",
    "## Install packages\n",
    "\n",
    "To use the new functionality, ensure the required packages are installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install geoai-py timm segmentation-models-pytorch lightning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import geoai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explore Available Encoders\n",
    "\n",
    "The timm library provides 1000+ encoders that can be used with segmentation architectures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List some popular encoders\n",
    "print(\"ResNet encoders:\", geoai.list_timm_models(filter=\"resnet\", limit=5))\n",
    "print(\"EfficientNet encoders:\", geoai.list_timm_models(filter=\"efficientnet\", limit=5))\n",
    "print(\"ConvNeXt encoders:\", geoai.list_timm_models(filter=\"convnext\", limit=5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Sample Data\n",
    "\n",
    "We'll use the same NAIP building detection dataset as the `train_segmentation_model` example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download NAIP aerial imagery and building footprints\n",
    "train_raster_url = (\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_rgb_train.tif\"\n",
    ")\n",
    "train_vector_url = \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_train_buildings.geojson\"\n",
    "test_raster_url = (\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_test.tif\"\n",
    ")\n",
    "\n",
    "train_raster_path = geoai.download_file(train_raster_url)\n",
    "train_vector_path = geoai.download_file(train_vector_url)\n",
    "test_raster_path = geoai.download_file(test_raster_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Sample Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoai.view_vector_interactive(train_vector_path, tiles=train_raster_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Training Data\n",
    "\n",
    "Generate image chips and corresponding segmentation masks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_folder = \"timm_buildings\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create training chips\n",
    "tiles = geoai.export_geotiff_tiles(\n",
    "    in_raster=train_raster_path,\n",
    "    out_folder=out_folder,\n",
    "    in_class_data=train_vector_path,\n",
    "    tile_size=512,\n",
    "    stride=256,\n",
    "    buffer_radius=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train U-Net with ResNet50 Encoder\n",
    "\n",
    "The `train_timm_segmentation_model` function provides a simplified interface similar to `train_segmentation_model`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train U-Net with ResNet50 encoder\n",
    "geoai.train_timm_segmentation_model(\n",
    "    images_dir=f\"{out_folder}/images\",\n",
    "    labels_dir=f\"{out_folder}/labels\",\n",
    "    output_dir=f\"{out_folder}/unet_resnet50\",\n",
    "    encoder_name=\"resnet50\",\n",
    "    architecture=\"unet\",\n",
    "    encoder_weights=\"imagenet\",\n",
    "    num_channels=3,\n",
    "    num_classes=2,  # background and building\n",
    "    batch_size=8,\n",
    "    num_epochs=20,\n",
    "    learning_rate=0.001,\n",
    "    val_split=0.2,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train DeepLabV3+ with EfficientNet-B3 Encoder\n",
    "\n",
    "EfficientNet encoders provide excellent performance with fewer parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train DeepLabV3+ with EfficientNet-B3\n",
    "geoai.train_timm_segmentation_model(\n",
    "    images_dir=f\"{out_folder}/images\",\n",
    "    labels_dir=f\"{out_folder}/labels\",\n",
    "    output_dir=f\"{out_folder}/deeplabv3plus_efficientnet_b3\",\n",
    "    encoder_name=\"efficientnet-b3\",  # Note: use dash for SMP encoders\n",
    "    architecture=\"deeplabv3plus\",\n",
    "    encoder_weights=\"imagenet\",\n",
    "    num_channels=3,\n",
    "    num_classes=2,\n",
    "    batch_size=8,\n",
    "    num_epochs=20,\n",
    "    learning_rate=0.001,\n",
    "    val_split=0.2,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Performance Analysis\n",
    "\n",
    "Let's examine the training curves and model performance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoai.plot_performance_metrics(\n",
    "    history_path=f\"{out_folder}/unet_resnet50/models/training_history.pth\",\n",
    "    figsize=(15, 5),\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance Metrics\n",
    "\n",
    "**IoU (Intersection over Union)** is the primary metric used to evaluate semantic segmentation performance.\n",
    "\n",
    "### 🔸 **IoU Definition**\n",
    "\n",
    "$$\n",
    "\\text{IoU} = \\frac{|A \\cap B|}{|A \\cup B|} = \\frac{TP}{TP + FP + FN}\n",
    "$$\n",
    "\n",
    "- Measures the overlap between predicted region $A$ and ground truth region $B$ relative to their union\n",
    "- Ranges from 0 (no overlap) to 1 (perfect overlap)\n",
    "- Common in object detection and semantic segmentation benchmarks (e.g., COCO, Pascal VOC)\n",
    "\n",
    "The training curves show:\n",
    "- **Training Loss**: How well the model fits the training data\n",
    "- **Validation Loss**: How well the model generalizes to unseen data\n",
    "- **Validation IoU**: The overlap accuracy on validation data\n",
    "\n",
    "> **Note**: Higher IoU is better (closer to 1.0), lower loss is better (closer to 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fine-tune with frozen encoder\n",
    "geoai.train_timm_segmentation_model(\n",
    "    images_dir=f\"{out_folder}/images\",\n",
    "    labels_dir=f\"{out_folder}/labels\",\n",
    "    output_dir=f\"{out_folder}/unet_resnet50_frozen\",\n",
    "    encoder_name=\"resnet50\",\n",
    "    architecture=\"unet\",\n",
    "    encoder_weights=\"imagenet\",\n",
    "    num_channels=3,\n",
    "    num_classes=2,\n",
    "    freeze_encoder=True,  # Freeze encoder weights\n",
    "    batch_size=8,\n",
    "    num_epochs=10,  # Fewer epochs needed\n",
    "    learning_rate=0.001,\n",
    "    val_split=0.2,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Inference on Test Image\n",
    "\n",
    "Use the trained model to segment the test image using the `timm_semantic_segmentation` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inference\n",
    "masks_path = \"naip_test_timm_prediction.tif\"\n",
    "model_path = f\"{out_folder}/unet_resnet50/models/last.ckpt\"\n",
    "\n",
    "geoai.timm_semantic_segmentation(\n",
    "    input_path=test_raster_path,\n",
    "    output_path=masks_path,\n",
    "    model_path=model_path,\n",
    "    encoder_name=\"resnet50\",\n",
    "    architecture=\"unet\",\n",
    "    num_channels=3,\n",
    "    num_classes=2,\n",
    "    window_size=512,\n",
    "    overlap=256,\n",
    "    batch_size=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vectorize and Visualize Results\n",
    "\n",
    "Convert the segmentation mask to vector format and visualize:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vectorize the mask\n",
    "output_vector_path = \"naip_test_timm_prediction.geojson\"\n",
    "gdf = geoai.orthogonalize(masks_path, output_vector_path, epsilon=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add geometric properties\n",
    "gdf_props = geoai.add_geometric_properties(gdf, area_unit=\"m2\", length_unit=\"m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize results\n",
    "geoai.view_raster(masks_path, nodata=0, basemap=test_raster_path, backend=\"ipyleaflet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter buildings by area and visualize\n",
    "gdf_filtered = gdf_props[gdf_props[\"area_m2\"] > 50]\n",
    "geoai.view_vector_interactive(gdf_filtered, column=\"area_m2\", tiles=test_raster_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create split map comparison\n",
    "geoai.create_split_map(\n",
    "    left_layer=gdf_filtered,\n",
    "    right_layer=test_raster_path,\n",
    "    left_args={\"style\": {\"color\": \"red\", \"fillOpacity\": 0.2}},\n",
    "    basemap=test_raster_path,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hugging Face Hub Integration\n",
    "\n",
    "The geoai library now supports loading models from and pushing models to the Hugging Face Hub. This enables:\n",
    "\n",
    "- **Loading Pre-trained Models**: Use state-of-the-art segmentation models from HF Hub\n",
    "- **Sharing Your Models**: Upload trained models to share with the community\n",
    "- **Model Versioning**: Leverage HF Hub's versioning and collaboration features\n",
    "  \n",
    "\n",
    "### Option 1: Load a Pre-trained Model from HF Hub\n",
    "\n",
    "You can use complete segmentation models from Hugging Face Hub with the `use_timm_model=True` parameter:\n",
    "\n",
    "**Note**: This example is commented out as it requires a specific HF Hub model. Uncomment if you have a compatible model available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Example: Load and fine-tune a model from HF Hub\n",
    "# geoai.train_timm_segmentation_model(\n",
    "#     images_dir=f\"{out_folder}/images\",\n",
    "#     labels_dir=f\"{out_folder}/labels\",\n",
    "#     output_dir=f\"{out_folder}/hf_hub_model\",\n",
    "#     use_timm_model=True,\n",
    "#     timm_model_name=\"hf-hub:username/model-name\",  # Replace with actual HF Hub model\n",
    "#     num_channels=3,\n",
    "#     num_classes=2,\n",
    "#     batch_size=8,\n",
    "#     num_epochs=10,\n",
    "#     learning_rate=0.0001,  # Lower LR for fine-tuning\n",
    "#     val_split=0.2,\n",
    "#     verbose=True,\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2: Load a Pushed Model for Inference\n",
    "\n",
    "Once a model is pushed to HF Hub, you can load it for inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Example: Load model from HF Hub and run inference\n",
    "# # First, download the model from HF Hub (this would happen automatically)\n",
    "# # Then run inference with the downloaded model\n",
    "# geoai.timm_semantic_segmentation(\n",
    "#     input_path=test_raster_path,\n",
    "#     output_path=\"naip_test_hf_hub_prediction.tif\",\n",
    "#     model_path=\"path/to/downloaded/model.pth\",  # Path to downloaded HF Hub model\n",
    "#     encoder_name=\"resnet50\",\n",
    "#     architecture=\"unet\",\n",
    "#     num_channels=3,\n",
    "#     num_classes=2,\n",
    "#     use_timm_model=False,  # Set to True if it's a pure timm model\n",
    "#     window_size=512,\n",
    "#     overlap=256,\n",
    "#     batch_size=4,\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 3: Push Your Trained Model to HF Hub\n",
    "\n",
    "After training a model, you can share it on Hugging Face Hub:\n",
    "\n",
    "**Note**: This requires you to be logged in to Hugging Face. Run `huggingface-cli login` in your terminal first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Example: Push your trained model to HF Hub\n",
    "# url = geoai.push_timm_model_to_hub(\n",
    "#     model_path=f\"{out_folder}/unet_resnet50/models/last.ckpt\",\n",
    "#     repo_id=\"your-username/building-segmentation-resnet50-unet\",  # Replace with your repo\n",
    "#     encoder_name=\"resnet50\",\n",
    "#     architecture=\"unet\",\n",
    "#     num_channels=3,\n",
    "#     num_classes=2,\n",
    "#     commit_message=\"Upload building segmentation model trained on NAIP imagery\",\n",
    "#     private=False,  # Set to True for private repository\n",
    "# )\n",
    "# print(f\"Model uploaded to: {url}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hugging Face Hub Integration\n",
    "\n",
    "The geoai library now supports loading models from and pushing models to the Hugging Face Hub. This enables:\n",
    "\n",
    "- **Loading Pre-trained Models**: Use state-of-the-art segmentation models from HF Hub\n",
    "- **Sharing Your Models**: Upload trained models to share with the community\n",
    "- **Model Versioning**: Leverage HF Hub's versioning and collaboration features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "\n",
    "1. **Simplified API**: Using `train_timm_segmentation_model()` similar to `train_segmentation_model()`\n",
    "2. **Multiple Encoders**: Training with ResNet50, EfficientNet-B3, and ConvNeXt-Tiny\n",
    "3. **Multiple Architectures**: U-Net, DeepLabV3+, and FPN\n",
    "4. **Transfer Learning**: Fine-tuning with frozen encoders\n",
    "5. **Inference**: Using `timm_semantic_segmentation()` for predictions\n",
    "6. **Post-processing**: Vectorization and visualization\n",
    "\n",
    "## Supported Architectures\n",
    "\n",
    "- **U-Net**: Classic encoder-decoder architecture\n",
    "- **U-Net++**: Nested U-Net with dense skip connections\n",
    "- **DeepLabV3**: Atrous Spatial Pyramid Pooling (ASPP)\n",
    "- **DeepLabV3+**: DeepLabV3 with decoder\n",
    "- **FPN**: Feature Pyramid Network\n",
    "- **PSPNet**: Pyramid Scene Parsing Network\n",
    "- **LinkNet**: Efficient architecture with skip connections\n",
    "- **MANet**: Multi-scale Attention Network\n",
    "- **PAN**: Pyramid Attention Network\n",
    "\n",
    "## Popular Encoders\n",
    "\n",
    "- **ResNet family**: resnet18, resnet34, resnet50, resnet101, resnet152\n",
    "- **EfficientNet family**: efficientnet_b0 to efficientnet_b7\n",
    "- **ConvNeXt family**: convnext_tiny, convnext_small, convnext_base\n",
    "- **RegNet family**: regnetx_002, regnetx_004, regnety_002, regnety_004\n",
    "- **MobileNet family**: mobilenetv2_100, mobilenetv3_large_100\n",
    "\n",
    "## Key Advantages\n",
    "\n",
    "1. **1000+ Encoders**: Access to state-of-the-art backbones from timm\n",
    "2. **Simple API**: Functions match the existing `train_segmentation_model` interface\n",
    "3. **Automatic Preprocessing**: Handles data loading and splitting automatically\n",
    "4. **Lightning Integration**: Built-in checkpointing, early stopping, and logging\n",
    "5. **IoU Monitoring**: Track IoU metrics during training\n",
    "\n",
    "## Next Steps\n",
    "\n",
    "- Experiment with different encoder-architecture combinations\n",
    "- Try modern encoders like ConvNeXt or Swin Transformer\n",
    "- Use data augmentation for better generalization\n",
    "- Apply to multi-class segmentation tasks"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
